from autogen import AssistantAgent, UserProxyAgent, GroupChat, GroupChatManager, Agent
import json
import re
import os
import autogen
import logging
from typing import Optional, List, Dict, Union, Callable, Any

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    filename='agent_interaction.log',
    encoding='UTF-8',
    filemode='w'
)
logger = logging.getLogger(__name__)

# 在控制台也显示日志
console = logging.StreamHandler()
console.setLevel(logging.INFO)
formatter = logging.Formatter('%(name)-12s: %(levelname)-8s %(message)s')
console.setFormatter(formatter)
logging.getLogger('').addHandler(console)


# 配置类
class Config:
    BASE_URL = "https://api.chatfire.cn/v1"
    API_KEY = "sk-3W5C4oPOVpaRjfihonOfPpPu7i96Wv9CS9RCdrYDhnmU7DAZ"
    MODEL = "gpt-4o"
    MAX_PROTOTYPE_ATTEMPTS = 5  # 原型设计师最大尝试次数
    QA_SCORE_THRESHOLD = 98  # 质量保证分数阈值


# 配置列表
config_list = [
    {
        "model": Config.MODEL,
        "api_key": Config.API_KEY,
        "base_url": Config.BASE_URL
    }
]


# 外部工具类
class ExternalTools:
    @staticmethod
    def save_prototype(code: str, filename: str) -> Dict[str, Any]:
        """保存原型代码到文件"""
        try:
            os.makedirs("prototypes", exist_ok=True)
            filepath = os.path.join("prototypes", filename)
            with open(filepath, 'w', encoding='utf-8') as f:
                f.write(code)
            logger.info(f"原型已保存至: {filepath}")
            return {"status": "success", "filepath": filepath}
        except Exception as e:
            logger.error(f"保存原型失败: {str(e)}")
            return {"status": "error", "message": str(e)}

    @staticmethod
    def generate_uml(requirements: Dict[str, Any]) -> Dict[str, Any]:
        """生成UML图（伪实现）"""
        try:
            logger.info(f"生成UML图: {requirements['system_name']}")
            # 实际项目中应集成PlantUML或Mermaid
            return {
                "status": "success",
                "details": f"UML图已生成: 包含{len(requirements['use_cases'])}个用例"
            }
        except KeyError as e:
            logger.error(f"UML生成失败: 缺少必要字段 {str(e)}")
            return {"status": "error", "message": f"缺少字段: {str(e)}"}


# 代理工厂类
class AgentFactory:
    @staticmethod
    def create_agents():
        # 用户代理（禁用代码执行）
        user_proxy = UserProxyAgent(
            name="User_Proxy",
            human_input_mode="NEVER",
            code_execution_config=False,  # 完全禁用代码执行
            function_map={
                "save_prototype": ExternalTools.save_prototype,
                "generate_uml": ExternalTools.generate_uml
            },
            max_consecutive_auto_reply=2  # 防止无限循环
        )

        # 需求分析师 - 加强系统提示
        requirement_analyst = AssistantAgent(
            name="Requirement_Analyst",
            system_message=(
                "你负责分析用户需求并生成结构化需求规格说明。"
                "输出必须是纯JSON格式，不要包含任何解释或额外文本。"
                "输出格式: {\"system_name\": \"名称\", \"core_function\": [\"功能1\", \"功能2\"], "
                "\"user_roles\": [\"角色1\", \"角色2\"], \"use_cases\": [{\"name\": \"用例名\", \"steps\": [\"步骤1\"]}]}"
            ),
            llm_config={"config_list": config_list}
        )

        # 原型设计师 - 加强系统提示
        prototype_designer = AssistantAgent(
            name="Prototype_Designer",
            system_message=(
                "你负责根据需求规格生成HTML原型代码。"
                "只输出完整的HTML代码，不要包含任何解释、注释或代码块标记(```)。"
                "使用Bootstrap框架创建响应式界面。"
                "确保代码可以直接在浏览器中运行。"
            ),
            llm_config={"config_list": config_list}
        )

        # 质量保证 - 加强系统提示（新增评分功能）
        quality_assurance = AssistantAgent(
            name="Quality_Assurance",
            system_message=(
                "你负责验证原型代码是否符合标准。"
                "在原型设计师生成原型后，你会自动被调用进行评审。"
                "检查完整性和功能性后，给出一个0-100分的评分。"
                "评分格式: 【评分:XX/100】，其中XX是具体分数。"
                "如果发现严重问题，直接指出具体问题。"
                "不要生成代码或修改内容，只需报告问题和给出评分。"
                "如果原型质量很高，请给出90分以上的评分。"
            ),
            llm_config={"config_list": config_list}
        )

        return user_proxy, requirement_analyst, prototype_designer, quality_assurance


# 提取原型代码 - 增强版
def extract_prototype_code(messages):
    for msg in messages:
        if msg['name'] == 'Prototype_Designer':
            content = msg.get('content', '')

            # 情况1: 内容已经是纯HTML
            if '<!DOCTYPE html>' in content:
                return content

            # 情况2: 包含代码块标记
            if '```html' in content:
                parts = content.split('```html')
                if len(parts) > 1:
                    return parts[1].split('```')[0].strip()

            # 情况3: 包含HTML标签
            if '<html>' in content or '<div>' in content:
                return content

    logger.warning("未找到原型代码")
    return ""


# 提取JSON需求 - 增强版
def extract_requirements(content):
    """从内容中提取JSON需求"""
    try:
        # 尝试直接解析
        return json.loads(content)
    except json.JSONDecodeError:
        # 尝试提取代码块中的JSON
        if '```json' in content:
            json_str = content.split('```json')[1].split('```')[0].strip()
            try:
                return json.loads(json_str)
            except:
                pass

        # 尝试提取可能的JSON部分
        json_match = re.search(r'\{[\s\S]*\}', content)
        if json_match:
            try:
                return json.loads(json_match.group(0))
            except:
                pass

    logger.error("无法提取有效的JSON需求")
    return None


# 提取质量保证评分 - 新增功能
def extract_qa_score(content):
    """从质量保证代理的回复中提取评分"""
    try:
        # 尝试匹配评分格式
        score_match = re.search(r'【评分[:：]\s*(\d+)/100】', content)
        if score_match:
            return int(score_match.group(1))

        # 尝试匹配纯数字分数
        score_match = re.search(r'\b(\d{1,3})/100\b', content)
        if score_match:
            return int(score_match.group(1))

        # 尝试匹配分数描述
        if "90分以上" in content or "优秀" in content or "完美" in content:
            return 95
        if "80分以上" in content or "良好" in content:
            return 85
        if "70分以上" in content or "合格" in content:
            return 75
        if "不及格" in content or "不合格" in content:
            return 60

        logger.warning("未找到明确评分，尝试估算")

        # 根据关键词估算分数
        positive_keywords = ["优秀", "完美", "非常好", "高质量", "符合标准"]
        negative_keywords = ["问题", "缺陷", "不足", "改进", "错误"]

        positive_count = sum(1 for kw in positive_keywords if kw in content)
        negative_count = sum(1 for kw in negative_keywords if kw in content)

        # 基础分 + 正面关键词加分 - 负面关键词减分
        base_score = 70
        score = base_score + (positive_count * 5) - (negative_count * 5)
        return max(0, min(100, score))

    except Exception as e:
        logger.error(f"评分提取失败: {str(e)}")
        return 0


# 终止条件监控类 - 完全重构
class TerminationMonitor:
    def __init__(self):
        self.prototype_attempts = 0
        self.highest_score = 0
        self.terminated = False
        self.termination_reason = ""
        logger.info("创建终止条件监控器")

    def increment_prototype_attempts(self):
        """增加原型设计尝试次数"""
        self.prototype_attempts += 1
        logger.info(f"原型设计尝试次数: {self.prototype_attempts}/{Config.MAX_PROTOTYPE_ATTEMPTS}")

        # 检查是否达到最大尝试次数
        if self.prototype_attempts >= Config.MAX_PROTOTYPE_ATTEMPTS:
            self.terminated = True
            self.termination_reason = f"达到最大原型设计尝试次数 ({Config.MAX_PROTOTYPE_ATTEMPTS})"
            logger.warning(self.termination_reason)

    def update_qa_score(self, score):
        """更新质量保证评分"""
        if score > self.highest_score:
            self.highest_score = score
            logger.info(f"更新最高质量评分: {self.highest_score}/100")

            # 检查是否达到分数阈值
            if score >= Config.QA_SCORE_THRESHOLD:
                self.terminated = True
                self.termination_reason = f"质量评分达到阈值 ({score}/{Config.QA_SCORE_THRESHOLD}+)"
                logger.warning(self.termination_reason)

    def should_terminate(self):
        """检查是否应该终止流程"""
        return self.terminated

    def get_status(self):
        """获取当前状态信息"""
        return {
            "prototype_attempts": self.prototype_attempts,
            "highest_score": self.highest_score,
            "terminated": self.terminated,
            "reason": self.termination_reason
        }


# 主工作流 - 使用新的终止条件
def run_workflow(user_input):
    logger.info("启动工作流")

    # 初始化代理
    user_proxy, analyst, designer, qa = AgentFactory.create_agents()

    # 创建终止条件监控器
    termination_monitor = TerminationMonitor()

    # 自定义群组聊天管理器 - 使用新的终止条件
    class SmartGroupChatManager(GroupChatManager):
        def generate_reply(
                self,
                messages: Optional[List[Dict]] = None,
                sender: Optional[Agent] = None,
                **kwargs
        ) -> Union[str, Dict, None]:
            # 检查终止条件
            if termination_monitor.should_terminate():
                logger.warning(f"终止回复生成: {termination_monitor.termination_reason}")
                return None

            # 调用原始生成方法
            reply = super().generate_reply(messages, sender, **kwargs)

            # 处理代理回复
            if sender and reply:
                # 记录原型设计师的尝试
                if sender.name == "Prototype_Designer":
                    termination_monitor.increment_prototype_attempts()

                # 处理质量保证代理的评分
                if sender.name == "Quality_Assurance":
                    try:
                        # 尝试提取评分
                        qa_score = extract_qa_score(reply)
                        termination_monitor.update_qa_score(qa_score)

                        # 在回复中添加评分信息
                        if qa_score > 0:
                            return f"{reply}\n\n【系统记录评分: {qa_score}/100】"
                    except Exception as e:
                        logger.error(f"处理质量评分时出错: {str(e)}")

            return reply

    # 创建群组聊天
    group_chat = GroupChat(
        agents=[user_proxy, analyst, designer, qa],
        messages=[],
        max_round=20,  # 设置足够大的上限作为安全网
        speaker_selection_method="auto",
        allow_repeat_speaker=False
    )

    # 使用自定义管理器
    manager = SmartGroupChatManager(
        groupchat=group_chat,
        llm_config={"config_list": config_list}
    )

    # 启动需求分析
    logger.info("启动需求分析...")
    try:
        user_proxy.initiate_chat(
            manager,
            message=f"用户需求: {user_input}\n请生成结构化需求规格说明(JSON格式)。"
        )
    except Exception as e:
        logger.error(f"聊天初始化失败: {str(e)}")
        return {"status": "error", "message": f"聊天初始化失败: {str(e)}"}

    # 检查是否因终止条件提前终止
    if termination_monitor.should_terminate():
        logger.warning(f"流程提前终止: {termination_monitor.termination_reason}")
        return {
            "status": "partial",
            "message": termination_monitor.termination_reason,
            "completed_steps": "需求分析"
        }

    # 提取需求规格
    requirements = None
    for msg in group_chat.messages:
        if msg['name'] == 'Requirement_Analyst':
            requirements = extract_requirements(msg.get('content', ''))
            if requirements:
                break

    if not requirements:
        logger.error("需求分析失败")
        return {"status": "error", "message": "需求分析失败"}

    logger.info(f"需求分析成功: {json.dumps(requirements, indent=2, ensure_ascii=False)}")

    # 检查终止条件
    if termination_monitor.should_terminate():
        return {
            "status": "partial",
            "message": termination_monitor.termination_reason,
            "completed_steps": "需求分析",
            "requirements": requirements
        }

    # 生成UML图
    logger.info("生成UML图...")
    try:
        uml_result = user_proxy.execute_function(
            func_call={
                "name": "generate_uml",
                "arguments": json.dumps({"requirements": requirements}, ensure_ascii=False)
            }
        )
        logger.info(f"UML生成结果: {uml_result}")
    except Exception as e:
        logger.error(f"UML生成失败: {str(e)}")
        uml_result = {"error": str(e)}

    # 检查终止条件
    if termination_monitor.should_terminate():
        return {
            "status": "partial",
            "message": termination_monitor.termination_reason,
            "completed_steps": "需求分析和UML生成",
            "requirements": requirements,
            "uml": uml_result
        }

    # 原型生成与质量保证循环
    logger.info("开始原型生成与质量保证循环...")
    prototype_code = ""
    for attempt in range(1, Config.MAX_PROTOTYPE_ATTEMPTS + 1):
        if termination_monitor.should_terminate():
            logger.info(f"终止条件已满足，跳过第 {attempt} 次尝试")
            break

        logger.info(f"尝试 #{attempt} - 生成原型")

        # 生成原型
        try:
            user_proxy.initiate_chat(
                manager,
                message=(
                    f"基于以下需求生成HTML原型(尝试 #{attempt}):\n{json.dumps(requirements, indent=2, ensure_ascii=False)}\n"
                    "重要提示：只输出HTML代码，不要包含任何解释、注释或代码块标记(```)。"
                    "确保代码可以直接在浏览器中运行。"
                    "生成后，请质量保证代理进行评审。"
                ),
                clear_history=False  # 保留历史对话
            )
        except Exception as e:
            logger.error(f"原型生成聊天失败: {str(e)}")
            continue

        # 提取原型代码
        current_prototype = extract_prototype_code(group_chat.messages)
        if current_prototype:
            prototype_code = current_prototype
            logger.info(f"成功提取第 {attempt} 次尝试的原型代码")

            # 请求质量保证评分
            logger.info(f"尝试 #{attempt} - 请求质量保证评审")
            try:
                user_proxy.initiate_chat(
                    manager,
                    message=(
                        f"请质量保证代理对最新生成的原型进行评审。\n"
                        f"这是第 {attempt} 次尝试的原型，请给出评分和改进建议。"
                    ),
                    clear_history=False  # 保留历史对话
                )
            except Exception as e:
                logger.error(f"质量保证请求失败: {str(e)}")
        else:
            logger.warning(f"第 {attempt} 次尝试未提取到有效的原型代码")

        # 检查终止条件（在质量保证评分后）
        if termination_monitor.should_terminate():
            logger.info(f"终止条件在尝试 #{attempt} 后满足")
            break

    # 保存最终原型
    if prototype_code:
        try:
            logger.info("保存最终原型...")
            save_result = user_proxy.execute_function(
                func_call={
                    "name": "save_prototype",
                    "arguments": json.dumps({
                        "code": prototype_code,
                        "filename": "generated_prototype.html"
                    }, ensure_ascii=False)
                }
            )
            logger.info(f"原型保存结果: {save_result}")

            # 返回成功结果
            return {
                "status": "success",
                "requirements": requirements,
                "uml": uml_result,
                "prototype": "prototypes/generated_prototype.html",
                "log_file": "agent_interaction.log",
                "qa_score": termination_monitor.highest_score,
                "prototype_attempts": termination_monitor.prototype_attempts,
                "termination_reason": termination_monitor.termination_reason
            }
        except Exception as e:
            logger.error(f"原型保存失败: {str(e)}")
            return {"status": "error", "message": f"原型保存失败: {str(e)}"}
    else:
        logger.error("未提取到有效的原型代码")
        return {"status": "error", "message": "未提取到有效的原型代码"}


# 主函数
if __name__ == "__main__":
    print(f"= 启动智能原型生成系统 =")
    print(f"AutoGen版本: {autogen.__version__}")
    print(f"配置参数: 最大原型尝试次数 = {Config.MAX_PROTOTYPE_ATTEMPTS}")
    print(f"配置参数: 质量分数阈值 = {Config.QA_SCORE_THRESHOLD}")

    # 示例需求
    user_input = "开发一个2048数字合并的游戏，界面要美观大方，简单显示游戏规则，重点是游戏一定要能玩，确保实现了游戏逻辑，要有开始游戏的按键，计分功能。"

    # 执行工作流
    result = run_workflow(user_input)

    print("\n" + "=" * 50)
    print("最终结果:")
    print(json.dumps(result, indent=2, ensure_ascii=False))

    # 结果状态处理
    status = result.get("status")
    if status == "success":
        print(f"\n✅ 原型生成成功！")
        print(f"📊 质量评分: {result.get('qa_score', '未知')}/100")
        print(f"🔄 原型尝试次数: {result.get('prototype_attempts', '未知')}")
        print(f"📁 文件位置: {result['prototype']}")
        print(f"📝 详细日志: {result['log_file']}")
    elif status == "partial":
        print(f"\n⚠️ 流程提前终止: {result.get('termination_reason', '未知原因')}")
        print(f"📊 最高质量评分: {result.get('qa_score', '未知')}/100")
        print(f"🔄 原型尝试次数: {result.get('prototype_attempts', '未知')}")
        print(f"已完成步骤: {result.get('completed_steps', '未知')}")
        print("ℹ️ 部分结果已保存")
    else:
        print(f"\n❌ 生成失败: {result.get('message', '未知错误')}")
        print("⚠️ 请查看日志文件获取详细信息")